package com.gitee.tanmingjian.httptunnel.core.util;

import com.gitee.tanmingjian.httptunnel.core.bean.*;
import com.gitee.tanmingjian.httptunnel.core.exception.WriteException;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

public class TunnelProtocolUtils {

    /**
     * header:
     * <pre>
     * ┌──────────────┬─────────────┬───────────────┬───────────┐
     * │ 1111(4 byte) │ 204(2 byte) │ errno(4 byte) │ 0(6 byte) │
     * └──────────────┴─────────────┴───────────────┴───────────┘
     * </pre>
     *
     * @param errno 错误码
     */
    private static ByteArrayOutputStream initHeader(int errno) {
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        try {
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(1111));
            byteArrayOutputStream.write(DataPackageUtils.getShortBinary(204));
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(errno));
            byteArrayOutputStream.write(DataPackageUtils.getDummy(6));
        } catch (IOException e) {
            throw new WriteException(e);
        }
        return byteArrayOutputStream;
    }

    /**
     * 数据库连接错误:
     * <pre>
     * ┌───────────────────┬──────────────────────┐
     * │ [header(16 byte)] │ error message(block) │
     * └───────────────────┴──────────────────────┘
     * </pre>
     * <ul>
     *   <li>header : {@link TunnelProtocolUtils#initHeader}</li>
     * </ul>
     *
     * @param errno 错误码
     * @param message 错误信息
     * @return 数据包
     */
    public static byte[] buildErrorByte(int errno, String message) {
        ByteArrayOutputStream byteArrayOutputStream = initHeader(errno);
        try {
            byteArrayOutputStream.write(DataPackageUtils.getBlock(message));
        } catch (IOException e) {
            throw new WriteException(e);
        }
        return byteArrayOutputStream.toByteArray();
    }

    /**
     * 测试连接:
     * <pre>
     * ┌───────────────────┬─────────────┬──────────────┬────────────────┐
     * │ [header(16 byte)] │ host(block) │ proto(block) │ version(block) │
     * └───────────────────┴─────────────┴──────────────┴────────────────┘
     * </pre>
     * <ul>
     *   <li>header : {@link TunnelProtocolUtils#initHeader}</li>
     * </ul>
     *
     * @param connectionInfo 数据库连接信息
     * @return 数据包
     */
    public static byte[] buildConnInfoByte(ConnectionInfo connectionInfo) {
        ByteArrayOutputStream byteArrayOutputStream = initHeader(0);

        try {
            byteArrayOutputStream.write(DataPackageUtils.getBlock(connectionInfo.getHost()));
            byteArrayOutputStream.write(DataPackageUtils.getBlock(connectionInfo.getProto()));
            byteArrayOutputStream.write(DataPackageUtils.getBlock(connectionInfo.getVersion()));
        } catch (IOException e) {
            throw new WriteException(e);
        }

        return byteArrayOutputStream.toByteArray();
    }

    /**
     * 查询数据:
     * <pre>
     * ┌───────────────────┬──────────────────┬──────────────────┬─────┐
     * │ [header(16 byte)] │ [query result 1] │ [query result 2] │ ... │
     * └───────────────────┴──────────────────┴──────────────────┴─────┘
     * query result:
     * ┌───────────────────┬───────────────────────────────────────────────┬────────────────────────┐
     * │ result set header │ fields header and data / error message(block) │ has next flag (1 byte) │
     * └───────────────────┴───────────────────────────────────────────────┴────────────────────────┘
     * </pre>
     * <ul>
     *   <li>result set header : {@link TunnelProtocolUtils#writeResultSetHeader}</li>
     *   <li>fields header and data : {@link TunnelProtocolUtils#writeResultSetData}</li>
     * </ul>
     * @param list 数据查询结果
     * @return 数据包
     */
    public static byte[] buildQueryResultsByte(List<QueryResult> list) {
        ByteArrayOutputStream byteArrayOutputStream = initHeader(0);

        for (QueryResult queryResult : list) {
            writeResultSetHeader(byteArrayOutputStream, queryResult.getResultSetHeader());

            Object resultSetBody = queryResult.getResultSetBody();

            if (resultSetBody instanceof ResultSetData) {
                writeResultSetData(byteArrayOutputStream, (ResultSetData) resultSetBody);
            } else if (resultSetBody instanceof String) {
                writeResultSetMessage(byteArrayOutputStream, (String) resultSetBody);
            } else {
                writeResultSetMessage(byteArrayOutputStream, StringUtils.EMPTY);
            }
            byteArrayOutputStream.write(Optional.ofNullable(queryResult.getHasNext()).orElse(false) ? 1 : 0);
        }

        return byteArrayOutputStream.toByteArray();
    }

    /**
     * result set header:
     * <pre>
     * ┌───────────────┬──────────────────────┬───────────────────┬────────────────────┬──────────────────┬────────────┐
     * │ errno(4 byte) │ affect rows (4 byte) │ insert id(4 byte) │ num fields(4 byte) │ num rows(4 byte) │ 0(12 byte) │
     * └───────────────┴──────────────────────┴───────────────────┴────────────────────┴──────────────────┴────────────┘
     * </pre>
     * @param byteArrayOutputStream 数据包
     * @param resultSetHeader 查询结果头信息
     */
    private static void writeResultSetHeader(ByteArrayOutputStream byteArrayOutputStream, ResultSetHeader resultSetHeader) {
        try {
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(resultSetHeader.getErrno()));
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(resultSetHeader.getAffectRows()));
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(resultSetHeader.getInsertId()));
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(resultSetHeader.getNumFields()));
            byteArrayOutputStream.write(DataPackageUtils.getLongBinary(resultSetHeader.getNumRows()));
            byteArrayOutputStream.write(DataPackageUtils.getDummy(12));
        } catch (IOException e) {
            throw new WriteException(e);
        }
    }

    /**
     * fields header and data:
     * <pre>
     * ┌──────────────────────────────────────────────────────────────────────────────────────────────────────────────┬─────┐
     * │┌───────────────────┬───────────────────┬───────────────────────┬─────────────────────┬──────────────────────┐│     │
     * ││ field name(block) │ table name(block) │ MySQL type id(4 byte) │ field flags(4 byte) │ field length(4 byte) ││ ... │
     * │└───────────────────┴───────────────────┴───────────────────────┴─────────────────────┴──────────────────────┘│     │
     * ├──────────────────────────────┬─────┬─────────────────────────────────────────────────────────────────────────┴─────┘
     * │┌────────────────────────────┐│     │
     * ││ data string or 0xFF(block) ││ ... │
     * │└────────────────────────────┘│     │
     * └──────────────────────────────┴─────┘
     * </pre>
     *
     * @param byteArrayOutputStream 数据包
     * @param resultSetData 查询结果头信息
     */
    private static void writeResultSetData(ByteArrayOutputStream byteArrayOutputStream, ResultSetData resultSetData) {
        writeFieldsHeader(byteArrayOutputStream, resultSetData.getFieldHeaders());
        writeData(byteArrayOutputStream, resultSetData.getData());
    }

    /**
     * message
     * <pre>
     * ┌────────────────┐
     * │ message(block) │
     * └────────────────┘
     * </pre>
     *
     * @param byteArrayOutputStream 数据包
     * @param message 信息文本
     */
    private static void writeResultSetMessage(ByteArrayOutputStream byteArrayOutputStream, String message) {
        try {
            byteArrayOutputStream.write(DataPackageUtils.getBlock(message));
        } catch (IOException e) {
            throw new WriteException(e);
        }
    }

    /**
     * fields header
     * <pre>
     * ┌──────────────────────────────────────────────────────────────────────────────────────────────────────────────┬─────┐
     * │┌───────────────────┬───────────────────┬───────────────────────┬─────────────────────┬──────────────────────┐│     │
     * ││ field name(block) │ table name(block) │ MySQL type id(4 byte) │ field flags(4 byte) │ field length(4 byte) ││ ... │
     * │└───────────────────┴───────────────────┴───────────────────────┴─────────────────────┴──────────────────────┘│     │
     * └──────────────────────────────────────────────────────────────────────────────────────────────────────────────┴─────┘
     * </pre>
     *
     * @param byteArrayOutputStream 数据包
     * @param fieldHeaders 字段信息
     */
    private static void writeFieldsHeader(ByteArrayOutputStream byteArrayOutputStream, List<FieldHeader> fieldHeaders) {
        try {
            for (FieldHeader fieldHeader : fieldHeaders) {
                byteArrayOutputStream.write(DataPackageUtils.getBlock(fieldHeader.getName()));
                byteArrayOutputStream.write(DataPackageUtils.getBlock(Optional.ofNullable(fieldHeader.getTableName()).orElse(StringUtils.EMPTY)));
                byteArrayOutputStream.write(DataPackageUtils.getLongBinary(fieldHeader.getTypeId()));
                byteArrayOutputStream.write(DataPackageUtils.getLongBinary(fieldHeader.getFlags()));
                byteArrayOutputStream.write(DataPackageUtils.getLongBinary(fieldHeader.getLength()));
            }
        } catch (IOException e) {
            throw new WriteException(e);
        }
    }

    /**
     * data
     * <pre>
     * ┌──────────────────────────────┬─────┐
     * │┌────────────────────────────┐│     │
     * ││ data string or 0xFF(block) ││ ... │
     * │└────────────────────────────┘│     │
     * └──────────────────────────────┴─────┘
     * </pre>
     *
     * @param byteArrayOutputStream 数据包
     * @param data 数据
     */
    private static void writeData(ByteArrayOutputStream byteArrayOutputStream, List<List<String>> data) {
        try {
            for (List<String> row : data) {
                for (String cell : row) {
                    byteArrayOutputStream.write(DataPackageUtils.getBlock(cell));
                }
            }
        } catch (IOException e) {
            throw new WriteException(e);
        }
    }

    /**
     * 数据包解析为连接信息
     * {@link TunnelProtocolUtils#buildConnInfoByte}</li>
     */
    public static ConnectionInfo toConnectionInfo(ByteBuffer byteBuffer) {
        // 校验头部
        validHeader(byteBuffer);
        // dummy

        String host = DataPackageUtils.getBlockString(byteBuffer);
        String proto = DataPackageUtils.getBlockString(byteBuffer);
        String version = DataPackageUtils.getBlockString(byteBuffer);

        return new ConnectionInfo(host, proto, version);
    }

    /**
     * 数据包解析为查询结果
     * {@link TunnelProtocolUtils#buildQueryResultsByte}</li>
     */
    public static List<QueryResult> toQueryResultList(ByteBuffer byteBuffer) {
        validHeader(byteBuffer);

        List<QueryResult> result = new ArrayList<>();

        boolean hasNext = true;
        while (hasNext) {
            QueryResult queryResult = new QueryResult();
            result.add(queryResult);
            // result set header
            int errno = DataPackageUtils.getLong(byteBuffer);
            int affectRows = DataPackageUtils.getLong(byteBuffer);
            int insertId = DataPackageUtils.getLong(byteBuffer);
            int numFields = DataPackageUtils.getLong(byteBuffer);
            int numRows = DataPackageUtils.getLong(byteBuffer);
            ResultSetHeader resultSetHeader = new ResultSetHeader(errno, affectRows, insertId, numFields, numRows);
            queryResult.setResultSetHeader(resultSetHeader);
            DataPackageUtils.getDummy(byteBuffer, 12);

            // fields header and data / error message(block)
            if (numFields == 0) {
                String message = DataPackageUtils.getBlockString(byteBuffer);
                queryResult.setResultSetBody(message);
            } else {
                // fields header and data
                List<FieldHeader> fieldHeaders = new ArrayList<>();
                for (int i = 0; i < numFields; i++) {
                    FieldHeader fieldHeader = new FieldHeader();
                    fieldHeader.setName(DataPackageUtils.getBlockString(byteBuffer));
                    fieldHeader.setTableName(DataPackageUtils.getBlockString(byteBuffer));
                    fieldHeader.setTypeId(DataPackageUtils.getLong(byteBuffer));
                    fieldHeader.setFlags(DataPackageUtils.getLong(byteBuffer));
                    fieldHeader.setLength(DataPackageUtils.getLong(byteBuffer));
                    fieldHeaders.add(fieldHeader);
                }

                List<List<String>> data = new ArrayList<>();
                for (int i = 0; i < numRows; i++) {
                    List<String> row = new ArrayList<>();
                    for (int j = 0; j < numFields; j++) {
                        row.add(DataPackageUtils.getBlockString(byteBuffer));
                    }
                    data.add(row);
                }
                ResultSetData resultSetData = new ResultSetData(fieldHeaders, data);
                queryResult.setResultSetBody(resultSetData);
            }

            byte b = byteBuffer.get();
            hasNext = b != 0;
            queryResult.setHasNext(hasNext);
        }

        return result;
    }

    private static void validHeader(ByteBuffer byteBuffer) {
        // 校验头部
        if (DataPackageUtils.getLong(byteBuffer) != 1111) {
            throw new RuntimeException();
        }
        if (DataPackageUtils.getShort(byteBuffer) != 204) {
            throw new RuntimeException();
        }
        int errno = DataPackageUtils.getLong(byteBuffer);
        if (errno != 0) {
            String message = DataPackageUtils.getBlockString(byteBuffer);
            throw new RuntimeException(message);
        }
        // dummy
        DataPackageUtils.getDummy(byteBuffer, 6);
    }

}
